#include <stdio.h>
#include <cuda_runtime.h>

// 简单的矩阵乘法内核 - 每个线程计算一个输出元素
__global__ void matrixMultiply(float *A, float *B, float *C, 
                              int numARows, int numAColumns,
                              int numBRows, int numBColumns) {
    // 计算当前线程要处理的输出位置
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (row < numARows && col < numBColumns) {
        float sum = 0.0f;
        for (int k = 0; k < numAColumns; k++) {
            sum += A[row * numAColumns + k] * B[k * numBColumns + col];
        }
        C[row * numBColumns + col] = sum;
    }
}

// 检查CUDA错误
void checkCudaError(cudaError_t err, const char* msg) {
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA Error: %s: %s\n", msg, cudaGetErrorString(err));
        exit(EXIT_FAILURE);
    }
}

int main() {
    const int M = 1024;  // A的行数
    const int K = 512;   // A的列数，B的行数
    const int N = 2048;  // B的列数
    
    // 主机内存分配
    size_t sizeA = M * K * sizeof(float);
    size_t sizeB = K * N * sizeof(float);
    size_t sizeC = M * N * sizeof(float);
    
    float *h_A = (float*)malloc(sizeA);
    float *h_B = (float*)malloc(sizeB);
    float *h_C = (float*)malloc(sizeC);
    
    // 初始化输入矩阵
    for (int i = 0; i < M * K; i++) h_A[i] = 1.0f;
    for (int i = 0; i < K * N; i++) h_B[i] = 2.0f;
    
    // 设备内存分配
    float *d_A, *d_B, *d_C;
    checkCudaError(cudaMalloc(&d_A, sizeA), "cudaMalloc d_A");
    checkCudaError(cudaMalloc(&d_B, sizeB), "cudaMalloc d_B");
    checkCudaError(cudaMalloc(&d_C, sizeC), "cudaMalloc d_C");
    
    // 拷贝数据到设备
    checkCudaError(cudaMemcpy(d_A, h_A, sizeA, cudaMemcpyHostToDevice), "Copy A to device");
    checkCudaError(cudaMemcpy(d_B, h_B, sizeB, cudaMemcpyHostToDevice), "Copy B to device");
    
    // 定义网格和块维度
    dim3 blockSize(16, 16);
    dim3 gridSize((N + blockSize.x - 1) / blockSize.x, 
                  (M + blockSize.y - 1) / blockSize.y);
    
    // 创建CUDA事件用于计时
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    
    // 执行内核
    cudaEventRecord(start);
    matrixMultiply<<<gridSize, blockSize>>>(d_A, d_B, d_C, M, K, K, N);
    cudaEventRecord(stop);
    
    // 等待内核完成并检查错误
    checkCudaError(cudaDeviceSynchronize(), "Kernel execution");
    checkCudaError(cudaGetLastError(), "Kernel error");
    
    // 计算执行时间
    float milliseconds = 0;
    cudaEventElapsedTime(&milliseconds, start, stop);
    printf("矩阵乘法执行时间: %.3f ms\n", milliseconds);
    
    // 拷贝结果回主机
    checkCudaError(cudaMemcpy(h_C, d_C, sizeC, cudaMemcpyDeviceToHost), "Copy C to host");
    
    // 验证结果（检查第一个元素）
    printf("结果矩阵C[0][0] = %.2f (期望值: %.2f)\n", h_C[0], 1.0f * 2.0f * K);
    
    // 清理
    free(h_A);
    free(h_B);
    free(h_C);
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    cudaEventDestroy(start);
    cudaEventDestroy(stop);
    
    return 0;
}